# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import re

from coin.base.datetime_util import to_timestamp_int
from coin.base.param_util import to_tuple
from experimental.prophet.graph.util import (from_data_type_proto, to_data_type_proto)


def to_duration_int(duration):
  if isinstance(duration, int):
    return duration
  elif isinstance(duration, float):
    return round(duration)
  elif isinstance(duration, str):
    return from_duration_str(duration)
  elif isinstance(duration, datetime.timedelta):
    return (duration.days * 86400 * 10**9 + duration.seconds * 10**9
            + duration.microseconds * 10**3)
  else:
    raise ValueError(duration)


def _count_trailing_zero(x):
  assert isinstance(x, int)
  count = 0
  while x:
    if x % 10 != 0:
      break
    count += 1
    x //= 10
  return count


def to_duration_str(duration):
  duration = to_duration_int(duration)
  if duration == 0:
    return ''

  NS = 1
  US = 10**3
  MS = 10**6
  S = 10**9
  MIN = 60 * S
  HOUR = 60 * MIN
  DAY = 24 * HOUR

  if duration >= DAY:
    prefix = '%dd' % (duration // DAY)
    return prefix + to_duration_str(duration - (duration // DAY) * DAY)
  elif duration >= HOUR:
    prefix = '%dh' % (duration // HOUR)
    return prefix + to_duration_str(duration - (duration // HOUR) * HOUR)
  elif duration >= MIN:
    prefix = '%dm' % (duration // MIN)
    return prefix + to_duration_str(duration - (duration // MIN) * MIN)
  elif duration >= S:
    prefix = '%ds' % (duration // S)
    return prefix + to_duration_str(duration - (duration // S) * S)
  elif duration >= MS:
    prefix = '%dms' % (duration // MS)
    return prefix + to_duration_str(duration - (duration // MS) * MS)
  elif duration >= US:
    prefix = '%dus' % (duration // US)
    return prefix + to_duration_str(duration - (duration // US) * US)
  else:
    return '%dns' % duration


def from_duration_str(duration_str, _all_str=None):
  assert isinstance(duration_str, str)
  if not duration_str:
    return 0

  _all_str = _all_str or duration_str
  units = {
      'ns': 1,
      'us': 10**3,
      'ms': 10**6,
      's': 10**9,
      'm': 60 * 10**9,
      'h': 3600 * 10**9,
      'd': 86400 * 10**9,
  }

  matched = re.match('([0-9]+(\.[0-9]+)?)(ns|us|ms|s|m|h|d)(.*)', duration_str)
  if not matched:
    raise ValueError(_all_str)

  num_part = matched.group(1)
  unit = matched.group(3)
  rest = matched.group(4)
  return round(float(num_part) * units[unit]) + from_duration_str(rest, _all_str)
